# -*- coding: utf-8 -*-

"""
 (c) 2023 - Copyright CTyunOS Inc

 Authors:
   youyifeng <youyf2@chinatelecom.cn>

"""
import os
from optparse import OptionParser
from cve_ease import activate_session, Scraper
from cve_ease.models import NVD
from cve_ease.helper import one_instance_check, download_file, ungzfile
import logging
import json

logger = logging.getLogger('cve-ease')


def get_usage_str(usage):
    return usage + "\n(Specify the --help global option for a list of other help options)"


def handle_nvd(gconfig, db_session, args):
    """[info] NVD info"""

    # one instance
    lock = one_instance_check(gconfig.LOCK_FILE_PATH)
    if lock:
        raise Exception(f"Another cve-ease already running")

    usage = "usage: %prog nvd <options>"
    parser = OptionParser(usage=get_usage_str(usage))

    parser.add_option('-m', '--makecache', dest='makecache', action='store_true', default=False,
                      help='get nvd cache')
    parser.add_option('--cleancache', dest='cleancache', action='store_true', default=False,
                      help='clean all cache and recreate db')
    parser.add_option('-l', '--list', dest='list', action='store_true', default=False,
                      help='list all nvd info')
    parser.add_option('-t', '--total', dest='total', action='store_true', default=False,
                      help='get nvd info statistics')
    parser.add_option('-v', '--verbose', dest='verbose', action='store_true', default=False,
                      help='show verbose output')

    (options, args) = parser.parse_args(args)

    if options.makecache:
        nvd_makecahe(gconfig, db_session)

    elif options.cleancache:
        if os.path.exists(gconfig.NVD_CACHE_PATH):
            os.removedirs(gconfig.NVD_CACHE_PATH)
        os.makedirs(gconfig.NVD_CACHE_PATH, exist_ok=True)

        # recreate nvd table
        NVD.__table__.drop(db_session.get_engine())
        NVD.__table__.create(db_session.get_engine())

        print(" Clean %s done!" % gconfig.NVD_CACHE_PATH)

    elif options.total:
        session = activate_session(db_session, gconfig)
        if gconfig.debug:
            print(" * active sql session")
        nvdDatabaseList = session.query(NVD).all()
        if 0 == len(nvdDatabaseList):
            print(" WARNNING: no nvd cache found! you should run 'cve-ease nvd -m' to cache it.")
            return
        print("total NVD:", len(nvdDatabaseList))
        nvd_number_by_year = {}
        for nvd in nvdDatabaseList:
            nvd_year = nvd.cve_id[4:8]
            if nvd_year not in nvd_number_by_year:
                nvd_number_by_year[nvd_year] = 1
            else:
                nvd_number_by_year[nvd_year] += 1
        nvd_sorted = sorted(nvd_number_by_year.items(), key=lambda item: item, reverse=True)

        for nvd_by_year in nvd_sorted:
            print(" %s : %d" % (nvd_by_year[0], nvd_by_year[1]))

    elif options.list:
        session = activate_session(db_session, gconfig)
        if gconfig.debug:
            print(" * active sql session")
        nvdDatabaseList = session.query(NVD).all()
        if 0 == len(nvdDatabaseList):
            print(" WARNNING: no nvd cache found! you should run 'cve-ease nvd -m' to cache it.")
            return
        for nvd in nvdDatabaseList:
            nvd.format_output()
    else:
        parser.print_help()


def nvd_makecahe(gconfig, db_session):
    os.makedirs(gconfig.NVD_CACHE_PATH, exist_ok=True)
    import datetime
    year = datetime.datetime.today().year
    begin_year = int(gconfig.NVD_BEGIN_YEAR) if gconfig.NVD_BEGIN_YEAR else 2005
    need_to_download = [str(y) for y in range(begin_year, year + 1)] + ['Modified']
    for index, yr in enumerate(need_to_download):
        targetfilename = "nvdcve-1.1-%s.json.gz" % str(yr)
        download_url = os.path.join(gconfig.NVD_BASE_URL, targetfilename)
        save_to = os.path.join(gconfig.NVD_CACHE_PATH, targetfilename)
        save_ungzfile = save_to[:-3]

        if os.path.exists(save_to):
            print(save_to, " already downloaded!")
        else:
            print("download %s to %s" % (download_url, save_to))
            download_file(download_url, save_to)

        if not os.path.exists(save_ungzfile):
            print("  ungz file to", save_ungzfile)
            ungzfile(save_to, save_ungzfile)
        print("[%5d/%-5d] Cache OK! %s" % (index + 1, len(need_to_download) - index - 1, save_to))

    print("download from %s done!" % gconfig.NVD_BASE_URL)

    session = activate_session(db_session, gconfig)
    if gconfig.debug:
        print(" * active sql session")

    nvdDatabaseList = []
    for index, yr in enumerate(need_to_download):
        targetfilename = "nvdcve-1.1-%s.json" % str(yr)
        stored_in = os.path.join(gconfig.NVD_CACHE_PATH, targetfilename)
        nvdDatabaseList.extend(extrace_nvd_info(stored_in))
        print("write %s to db done![%s]" % (targetfilename, len(nvdDatabaseList)))

    # recreate nvd table
    NVD.__table__.drop(db_session.get_engine())
    NVD.__table__.create(db_session.get_engine())

    session.bulk_insert_mappings(NVD, nvdDatabaseList)
    session.commit()
    print(" * makecache done!")


def extrace_nvd_info(cvejson_path):
    nvdDatabaseList = []
    json_data = None

    with open(cvejson_path, 'r') as jn:
        json_data = json.loads(jn.read())

    if json_data and 'CVE_Items' in json_data and len(json_data['CVE_Items']) > 0:
        for pcve in json_data['CVE_Items']:
            nvdDatabaseList.append({
                'cve_id': pcve['cve']['CVE_data_meta']['ID'],
                'description': pcve['cve']['description']['description_data'][0]['value'],
                'cve': pcve['cve'],
                'configurations': pcve['cve'],
                'impact': pcve['cve'],
                'publishedDate': pcve['publishedDate'],
                'lastModifiedDate': pcve['lastModifiedDate'],
            })
    return nvdDatabaseList
